import os
import re
from bs4 import BeautifulSoup
import pandas as pd
import time

# 创建存储数据的目录
def ensure_directory_exists(directory):
    if not os.path.exists(directory):
        os.makedirs(directory)

# 从HTML文件中提取前100名大学排名数据
def extract_top100_from_html(html_file_path, target_count=100):
    print(f'正在从文件提取数据: {html_file_path}')
    print(f'目标提取前{target_count}名大学数据')
    
    # 读取HTML文件
    try:
        with open(html_file_path, 'r', encoding='utf-8') as f:
            html_content = f.read()
    except Exception as e:
        print(f'读取HTML文件失败: {e}')
        return []
    
    # 使用BeautifulSoup解析HTML
    soup = BeautifulSoup(html_content, 'html.parser')
    
    # 初始化数据列表
    university_data = []
    
    # 常见国家名称列表，用于从大学名称中分离
    country_patterns = [
        'United States', 'United Kingdom', 'China', 'Japan', 'Germany', 
        'France', 'Canada', 'Australia', 'Switzerland', 'South Korea',
        'Singapore', 'Netherlands', 'Sweden', 'Denmark', 'Norway',
        'Italy', 'Spain', 'Israel', 'Hong Kong', 'Taiwan',
        'Finland', 'Belgium', 'New Zealand', 'Brazil', 'India'
    ]
    
    # 方法1: 尝试查找排名表格
    ranking_table = soup.find('table')
    if ranking_table:
        print('找到排名表格')
        rows = ranking_table.find_all('tr')
        print(f'表格包含{len(rows)}行数据')
        
        count = 0
        for row in rows:
            if count >= target_count:
                break
                
            cells = row.find_all('td')
            if len(cells) >= 2:
                # 提取排名
                rank = cells[0].get_text(strip=True)
                
                # 提取大学名称和国家（它们合并在一起）
                full_name = cells[1].get_text(strip=True)
                university_name = full_name
                country = ''
                
                # 尝试分离大学名称和国家
                for pattern in country_patterns:
                    if pattern in full_name:
                        university_name = full_name.replace(pattern, '').strip()
                        country = pattern
                        break
                
                # 提取关键统计数据（从其他列或页面元素中获取）
                # 注意：这部分需要根据网站实际结构调整
                student_staff_ratio = ''
                international_students = ''
                male_female_ratio = ''
                
                # 尝试从更多列中提取数据
                if len(cells) >= 6:
                    # 假设第4列是师生比，第5列是国际学生比例，第6列是男女比例
                    # 这些索引可能需要根据实际HTML结构调整
                    student_staff_ratio = cells[3].get_text(strip=True) if len(cells) > 3 else ''
                    international_students = cells[4].get_text(strip=True) if len(cells) > 4 else ''
                    male_female_ratio = cells[5].get_text(strip=True) if len(cells) > 5 else ''
                
                # 清理数据
                university_name = re.sub(r'^[-=\s]+|[-=\s]+$', '', university_name)
                
                # 只添加有有效大学名称的数据
                if university_name and len(university_name) > 3:
                    university_data.append([
                        rank,
                        university_name,
                        country,
                        student_staff_ratio,
                        international_students,
                        male_female_ratio
                    ])
                    count += 1
    
    # 如果方法1没有提取到足够数据，尝试方法2: 查找包含排名的div元素
    if len(university_data) < min(20, target_count):
        print(f'表格方法只提取到{len(university_data)}条数据，尝试查找包含排名的div元素...')
        
        # 查找可能包含排名的元素
        ranking_divs = soup.select('.rankings-table-container, .ranking-item, .university-item, tr')
        print(f'找到{len(ranking_divs)}个可能包含排名的元素')
        
        count = len(university_data)  # 从之前的计数继续
        for div in ranking_divs:
            if count >= target_count:
                break
                
            # 提取所有文本并检查是否包含排名模式
            text = div.get_text(separator=' ', strip=True)
            rank_match = re.search(r'(\d+|=\d+)\s+', text)
            
            if rank_match and len(text) > 20:
                rank = rank_match.group(1)
                remaining_text = text[rank_match.end():].strip()
                
                # 尝试提取大学名称和国家
                university_name = remaining_text
                country = ''
                
                for pattern in country_patterns:
                    if pattern in remaining_text:
                        university_name = remaining_text.replace(pattern, '').strip()
                        country = pattern
                        break
                
                # 清理数据
                university_name = re.sub(r'^[-=\s]+|[-=\s]+$', '', university_name)
                
                # 检查是否已经添加过这所大学
                is_duplicate = False
                for item in university_data:
                    if item[1] == university_name:
                        is_duplicate = True
                        break
                
                if university_name and len(university_name) > 3 and not is_duplicate:
                    university_data.append([
                        rank,
                        university_name,
                        country,
                        '', '', ''  # 关键统计数据（备用方法中可能无法提取）
                    ])
                    count += 1
    
    # 如果仍然没有足够数据，尝试从其他可能的来源提取
    if len(university_data) < min(20, target_count):
        print('尝试从页面所有文本中提取大学排名信息...')
        
        # 获取所有文本行
        all_text = soup.get_text(separator='\n')
        lines = all_text.split('\n')
        
        count = len(university_data)
        for line in lines:
            if count >= target_count:
                break
                
            line = line.strip()
            if len(line) > 20:
                # 检查行是否以排名数字开头
                rank_match = re.match(r'(\d+|=\d+)\s+', line)
                if rank_match:
                    rank = rank_match.group(1)
                    remaining_text = line[rank_match.end():].strip()
                    
                    # 尝试提取大学名称和国家
                    university_name = remaining_text
                    country = ''
                    
                    for pattern in country_patterns:
                        if pattern in remaining_text:
                            university_name = remaining_text.replace(pattern, '').strip()
                            country = pattern
                            break
                    
                    # 清理数据
                    university_name = re.sub(r'^[-=\s]+|[-=\s]+$', '', university_name)
                    
                    # 检查是否已经添加过这所大学
                    is_duplicate = False
                    for item in university_data:
                        if item[1] == university_name:
                            is_duplicate = True
                            break
                    
                    if university_name and len(university_name) > 3 and not is_duplicate:
                        university_data.append([
                            rank,
                            university_name,
                            country,
                            '', '', ''  # 关键统计数据
                        ])
                        count += 1
    
    # 限制结果为前100名
    university_data = university_data[:target_count]
    
    # 对数据进行排序（尽量按排名）
    try:
        # 转换排名为数字以便排序
        for i in range(len(university_data)):
            # 移除等号并尝试转换为数字
            rank_str = university_data[i][0].replace('=', '')
            try:
                rank_num = int(rank_str)
            except ValueError:
                rank_num = float('inf')  # 无法转换的排名放在最后
            # 在每个数据项前添加排名数字用于排序
            university_data[i] = [rank_num] + university_data[i]
        
        # 按排名数字排序
        university_data.sort(key=lambda x: x[0])
        
        # 移除临时添加的排名数字
        for i in range(len(university_data)):
            university_data[i] = university_data[i][1:]
        
        print('数据已按排名排序')
    except Exception as e:
        print(f'排序数据时出错: {e}，保持原始顺序')
    
    return university_data

# 保存数据到Excel文件
def save_to_excel(data):
    if not data or len(data) == 0:
        print('没有数据可保存')
        return False
    
    # 确定文件路径
    output_dir = 'd:\\学习的文件\\大学排名\\data'
    ensure_directory_exists(output_dir)
    timestamp = int(time.time())
    excel_file_path = os.path.join(output_dir, f'universities_ranking_top100_{timestamp}.xlsx')
    
    try:
        # 规范化数据格式
        print(f'正在处理{len(data)}条数据并保存到Excel...')
        
        # 定义标准表头
        headers = [
            '排名 (Ranking)', 
            '大学名称 (University Name)', 
            '国家 (Country)',
            '师生比 (Student-to-Staff Ratio)',
            '国际学生数量 (Number of International Students)',
            '男女比例 (Male-to-Female Ratio)'
        ]
        
        # 过滤和清理数据
        filtered_data = []
        for row in data:
            # 确保每行有6个元素
            while len(row) < 6:
                row.append('')
            if len(row) > 6:
                row = row[:6]  # 截断过长的行
            
            # 清理每个字段
            for i in range(len(row)):
                if isinstance(row[i], str):
                    row[i] = row[i].strip()
                    row[i] = re.sub('<[^<]+?>', '', row[i])  # 移除HTML标签
            
            # 只添加有有效大学名称的数据
            if row[1] and len(row[1]) > 3:
                filtered_data.append(row)
        
        # 创建DataFrame
        df = pd.DataFrame(filtered_data, columns=headers)
        
        # 保存到Excel
        with pd.ExcelWriter(excel_file_path, engine='openpyxl') as writer:
            df.to_excel(writer, index=False)
        
        print(f'数据已成功保存到Excel文件: {excel_file_path}')
        print(f'成功保存了{len(filtered_data)}条大学排名数据')
        
        # 显示前10条数据
        print('前10条数据预览:')
        for i in range(min(10, len(filtered_data))):
            print(f'{i+1}. 排名: {filtered_data[i][0]}, 大学: {filtered_data[i][1]}, 国家: {filtered_data[i][2]}')
        
        return True
    except Exception as e:
        print(f'保存Excel文件失败: {e}')
        # 备选方案：保存为CSV
        csv_file_path = os.path.join(output_dir, f'universities_ranking_top100_{timestamp}.csv')
        try:
            with open(csv_file_path, 'w', newline='', encoding='utf-8') as f:
                f.write(','.join(headers) + '\n')
                for row in filtered_data:
                    csv_row = [str(cell).replace(',', ';') for cell in row]
                    f.write(','.join(csv_row) + '\n')
            print(f'数据已保存为备选CSV文件: {csv_file_path}')
            return True
        except Exception as csv_error:
            print(f'保存CSV文件也失败: {csv_error}')
            return False

# 主函数
def main():
    print('启动大学排名数据提取工具（前100名）...')
    
    # 指定HTML文件路径 - 使用最近成功提取过数据的文件
    html_file_path = 'd:\\学习的文件\\大学排名\\data\\page_1_1760090164.html'
    
    # 检查文件是否存在
    if not os.path.exists(html_file_path):
        print(f'错误: 找不到HTML文件 {html_file_path}')
        print('请检查文件路径是否正确，或修改脚本中的html_file_path变量')
        return
    
    # 从HTML文件中提取前100名大学数据
    university_data = extract_top100_from_html(html_file_path, target_count=100)
    
    # 保存数据到Excel
    if university_data and len(university_data) > 0:
        print(f'成功从HTML文件提取到{len(university_data)}条大学排名数据！')
        save_to_excel(university_data)
    else:
        print('未能从HTML文件提取到足够的大学排名数据')
        print('尝试使用其他HTML文件或调整提取逻辑')

if __name__ == '__main__':
    main()